

import numpy as np
import matplotlib.pyplot as plt
from astropy.coordinates import EarthLocation, SkyCoord, AltAz, ICRS
from astropy.time import Time
from astropy import units as u
# from astropy_healpix import HEALPix
from datetime import datetime, timezone, timedelta
import matplotlib.patches as patches
import random
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import os

'''
TBD： 如何以天顶为原点，坐标轴和刻度显示赤道坐标系，同时对GWAC的每个圆中以文字显示camName
'''
'''
用天顶为中心，且坐标轴以天顶为原点的天球坐标系
正确绘制地平30度以上的区域、SVOM和GWAC覆盖范围
用天顶为中心半径60度的圆代表地平30度以上的天区
'''
def plot_sky_coverage(gwac_obs, svom_obs, target_time, savePath):
    """绘制以天顶为中心的天区覆盖图，包括地平30度以上、SVOM和GWAC覆盖范围"""
    
    # 目标时间
    # target_time = Time(datetime(2025, 2, 27, 19, 32, tzinfo=timezone.utc))
    # target_time = Time(datetime(2025, 2, 27, 12, 32, tzinfo=timezone.utc))
    target_time = Time(target_time)
    
    # 兴隆观测站位置
    obs_location = EarthLocation.from_geodetic(117.575, 40.393, 960)
    
    # 计算天顶的赤道坐标
    altaz_zenith = AltAz(alt=90*u.deg, az=0*u.deg, 
                         obstime=target_time, location=obs_location)
    zenith_coord = altaz_zenith.transform_to(ICRS())
    
    # 视场半径参数
    svom_field_radius = 90 * u.deg
    gwac_field_radius = 9.3 * u.deg
    horizon_radius = 60 * u.deg  # 地平30度以上的天顶角半径

    # 创建绘图
    fig, ax = plt.subplots(figsize=(8, 5), subplot_kw={'projection': 'aitoff'})
    
    # 绘制地平30度以上的天区覆盖范围（圆形区域）
    horizon_patch = patches.Circle(
        (0, 0),  # 天顶为中心
        radius=horizon_radius.to(u.rad).value,
        color='blue', alpha=0.2, label='Horizon > 30°'
    )
    ax.add_patch(horizon_patch)

    # 添加SVOM覆盖区域
    svom_center = None
    for start, end, center, _ in svom_obs:
        if start <= target_time.datetime <= end:
            svom_center = center
            break
            
    if svom_center is None:
        raise ValueError("SVOM观测时间不匹配")

    # 计算SVOM中心相对于天顶的偏移
    svom_delta_ra = (svom_center.ra - zenith_coord.ra).wrap_at(180 * u.deg)
    svom_delta_dec = svom_center.dec - zenith_coord.dec
    
    svom_patch = patches.Circle(
        (svom_delta_ra.radian, svom_delta_dec.radian),
        radius=svom_field_radius.to(u.rad).value,
        color='red', alpha=0.3, label='SVOM Field',
        transform=ax.transData
    )
    ax.add_patch(svom_patch)

    # 添加GWAC覆盖区域
    for tgObs in gwac_obs:
        obs_time = tgObs['time']
        if not (obs_time <= target_time.datetime <= obs_time + timedelta(minutes=10)):
            continue
            
        for ii, (camName, gwac_field) in enumerate(tgObs['telescope_data'].items()):
            # print(camName)
            # 计算GWAC相对天顶的偏移
            gwac_delta_ra = (gwac_field.ra - zenith_coord.ra).wrap_at(180 * u.deg)
            gwac_delta_dec = gwac_field.dec - zenith_coord.dec

            gwac_patch = patches.Circle(
                (gwac_delta_ra.radian, gwac_delta_dec.radian),
                radius=gwac_field_radius.to(u.rad).value,
                color="green",
                alpha=0.3, label=f'GWAC Field',
                transform=ax.transData
            )
            ax.add_patch(gwac_patch)
            
            # 添加相机名称标签
            ax.text(gwac_delta_ra.radian, gwac_delta_dec.radian, "G"+camName,
                    transform=ax.transData, fontsize=6, ha='center', va='center',
                    color='black', bbox=dict(facecolor='white', alpha=0.7, 
                                           edgecolor='none', boxstyle='round,pad=0.1'))
        break

    # 设置图例和标题
    handles, labels = ax.get_legend_handles_labels()
    by_label = dict(zip(labels, handles))  # 去重
    ax.legend(by_label.values(), by_label.keys(), loc='upper right')
    ax.set_title(f"sky coverage\n{target_time.iso}", pad=20)
    ax.grid(True)
    # plt.show()
    plt.savefig(savePath, dpi=300)
    plt.close('all')

'''
正确绘制SVOM和GWAC覆盖范围，地平30度以上的区域绘制错误
用Healpix均匀采样点代表地平30度以上的天区
'''
def plot_sky_coverage2(gwac_obs, svom_obs):
    """绘制以天顶为中心的天区覆盖图，包括地平30度以上、SVOM和GWAC覆盖范围"""
    
    # 兴隆观测站位置
    obs_location = EarthLocation.from_geodetic(117.575, 40.393, 960)
    
    # 目标时间
    target_time = Time(datetime(2025, 2, 27, 19, 32, tzinfo=timezone.utc))
    
    # 计算天顶的赤道坐标
    # altaz_zenith = AltAz(alt=90*u.deg, az=0*u.deg, 
    #                     obstime=target_time, location=obs_location)
    # zenith_coord = altaz_zenith.transform_to('icrs')
    from astropy.coordinates import ICRS  # 导入 ICRS 坐标框架

    # 计算天顶的赤道坐标
    altaz_zenith = AltAz(alt=90*u.deg, az=0*u.deg, 
                        obstime=target_time, location=obs_location)
    zenith_coord = altaz_zenith.transform_to(ICRS())  # 使用 ICRS() 而不是字符串 'icrs'
    
    # 视场半径参数
    svom_field_radius = 90 * u.deg
    gwac_field_radius = 9.3 * u.deg
    min_altitude = 30 * u.deg

    # 生成HEALPix网格
    nside = 32
    hp = HEALPix(nside=nside)
    ra, dec = hp.healpix_to_lonlat(np.arange(hp.npix))
    coords = SkyCoord(ra, dec, unit='deg')
    
    # 转换到AltAz坐标系筛选可见区域
    frame = AltAz(obstime=target_time, location=obs_location)
    altaz = coords.transform_to(frame)
    visible = altaz.alt > min_altitude

    # 计算所有点相对于天顶的坐标偏移
    delta_ra = (coords.ra - zenith_coord.ra).wrap_at(180 * u.deg)
    delta_dec = coords.dec - zenith_coord.dec

    # 创建Aitoff投影图
    fig, ax = plt.subplots(figsize=(10, 6), subplot_kw={'projection': 'aitoff'})
    
    # 绘制地平30度以上的可见区域
    ax.scatter(delta_ra[visible].radian, delta_dec[visible].radian,
               color='blue', s=1, alpha=0.5, label='Horizon > 30°')

    # 添加SVOM覆盖区域
    svom_center = None
    for start, end, center, _ in svom_obs:
        if start <= target_time.datetime <= end:
            svom_center = center
            break
            
    if svom_center is None:
        raise ValueError("SVOM观测时间不匹配")

    # 计算SVOM中心相对于天顶的偏移
    svom_delta_ra = (svom_center.ra - zenith_coord.ra).wrap_at(180 * u.deg)
    svom_delta_dec = svom_center.dec - zenith_coord.dec
    
    svom_patch = patches.Circle(
        (svom_delta_ra.radian, svom_delta_dec.radian),
        radius=svom_field_radius.to(u.rad).value,
        color='red', alpha=0.3, label='SVOM Field',
        transform=ax.transData
    )
    ax.add_patch(svom_patch)

    # 添加GWAC覆盖区域
    for tgObs in gwac_obs:
        obs_time = tgObs['time']
        if not (obs_time <= target_time.datetime <= obs_time + timedelta(minutes=10)):
            continue
            
        for ii, (_, gwac_field) in enumerate(tgObs['telescope_data'].items()):
            # 计算GWAC相对天顶的偏移
            gwac_delta_ra = (gwac_field.ra - zenith_coord.ra).wrap_at(180 * u.deg)
            gwac_delta_dec = gwac_field.dec - zenith_coord.dec

            gwac_patch = patches.Circle(
                (gwac_delta_ra.radian, gwac_delta_dec.radian),
                radius=gwac_field_radius.to(u.rad).value,
                color="green",
                alpha=0.3, label='GWAC Field',
                transform=ax.transData
            )
            ax.add_patch(gwac_patch)
            
        break

    # 设置图例和标题
    # ax.legend(loc='upper right', bbox_to_anchor=(1.1, 1))
    
    handles, labels = ax.get_legend_handles_labels()
    by_label = dict(zip(labels, handles))  # 去重
    ax.legend(by_label.values(), by_label.keys(), loc='upper right')
    
    ax.set_title(f" sky center \n{target_time.iso}", pad=20)
    ax.grid(True)
    plt.show()


'''
正确绘制SVOM覆盖范围，地平30度以上的区域绘制错误，GWAC的相机阵区域缺失多个
用Healpix均匀采样点代表地平30度以上的天区
'''
def plot_sky_coverage3(gwac_obs, svom_obs):
    """绘制 SVOM 天区范围、地平 30 度以上天区覆盖范围以及 GWAC 每个转台的天区覆盖范围"""
    
    # 兴隆观测站位置
    obs_location = EarthLocation.from_geodetic(117.575, 40.393, 960)
    
    # 目标时间
    target_time = Time(datetime(2025, 2, 27, 19, 32, tzinfo=timezone.utc))
    
    # SVOM 视场半径和 GWAC 视场半径
    svom_field_radius = 90 * u.deg
    gwac_field_radius = 9.3 * u.deg
    
    # 地平 30 度以上的天顶角
    min_altitude = 30 * u.deg
    cap_above_horizon_radius = 60 * u.deg  # 天顶角 = 90 - 高度角
    
    # 使用 HEALPix 网格生成天区坐标
    nside = 32
    hp = HEALPix(nside=nside)
    ra, dec = hp.healpix_to_lonlat(np.arange(hp.npix))
    coords = SkyCoord(ra, dec, unit='deg')
    
    # 转换到 AltAz 坐标系
    frame = AltAz(obstime=target_time, location=obs_location)
    altaz = coords.transform_to(frame)
    
    # 筛选地平 30 度以上的区域
    visible = altaz.alt > min_altitude
    
    # 获取 SVOM 的中心坐标
    svom_center = None
    for start, end, center, _ in svom_obs:
        if start <= target_time.datetime <= end:
            svom_center = center
            break
    
    if svom_center is None:
        raise ValueError("No SVOM observation found for the target time.")
    
    # 创建绘图
    fig, ax = plt.subplots(figsize=(10, 6), subplot_kw={'projection': 'aitoff'})
    
    # 将投影中心设置为 SVOM 的中心
    svom_ra = svom_center.ra.wrap_at(180 * u.deg).radian
    svom_dec = svom_center.dec.radian
    
    # 绘制地平 30 度以上的天区覆盖范围
    # visible_coords = coords[visible]
    # ax.scatter((visible_coords.ra.wrap_at(180 * u.deg) - svom_ra, 
    #            visible_coords.dec.radian - svom_dec, 
    #            color='blue', s=1, alpha=0.5, label='Horizon > 30°')
    # 绘制地平 30 度以上的天区覆盖范围
    visible_coords = coords[visible]
    ax.scatter(visible_coords.ra.wrap_at(180 * u.deg).radian - svom_ra, 
            visible_coords.dec.radian - svom_dec, 
            color='blue', s=1, alpha=0.5, label='Horizon > 30°')
    
    # 绘制 SVOM 的天区范围
    svom_patch = patches.Circle((0, 0),  # 中心已经设置为 SVOM 的中心
                               radius=svom_field_radius.to(u.rad).value, 
                               color='red', alpha=0.3, label='SVOM Field')
    ax.add_patch(svom_patch)
    
    # 绘制 GWAC 每个转台的天区覆盖范围
    for tgObs in gwac_obs:
        time = tgObs['time']
        telescopes = tgObs['telescope_data']
        # print(f"Time: {time}, Number of telescopes: {len(telescopes)}")
        
        # 检查目标时间是否在观测时间内
        if time <= target_time.datetime <= time + timedelta(minutes=10):  # 假设每个观测持续 10 分钟
            for ii, (_, gwac_field) in enumerate(telescopes.items()):
                # print(f"Telescope {ii + 1}: RA={gwac_field.ra.deg}, Dec={gwac_field.dec.deg}")
                
                # 计算 GWAC 中心相对于 SVOM 中心的偏移
                gwac_ra = gwac_field.ra.wrap_at(180 * u.deg).radian - svom_ra
                gwac_dec = gwac_field.dec.radian - svom_dec
                
                # 生成随机颜色
                random_color = (random.random(), random.random(), random.random())
                
                # 创建 Circle
                gwac_patch = patches.Circle((gwac_ra, gwac_dec), 
                                        radius=gwac_field_radius.to(u.rad).value, 
                                        color="green", alpha=0.3, label='GWAC Field',
                                        transform=ax.transData)  # 使用 Aitoff 投影的坐标变换
                
                # 添加 Circle 到图中
                ax.add_patch(gwac_patch)
            break
    
    # 设置图例和标题
    handles, labels = ax.get_legend_handles_labels()
    by_label = dict(zip(labels, handles))  # 去重
    ax.legend(by_label.values(), by_label.keys(), loc='upper right')
    
    ax.set_title(f"Sky Coverage at {target_time.iso} (Centered on SVOM)", pad=20)
    ax.grid(True)
    plt.show()
    
'''
使用cartopy的Mollweide球面投影绘制天区范围
是否正确有待验证
'''
def plot_sky_coverage4(gwac_obs, svom_obs):
    """绘制 SVOM 天区范围、地平 30 度以上天区覆盖范围以及 GWAC 每个转台的天区覆盖范围"""
    
    # 兴隆观测站位置
    obs_location = EarthLocation.from_geodetic(117.575, 40.393, 960)
    
    # 目标时间
    target_time = Time(datetime(2025, 2, 27, 19, 32, tzinfo=timezone.utc))
    
    # SVOM 视场半径和 GWAC 视场半径（以度为单位）
    svom_field_radius = 90  # 90 度
    gwac_field_radius = 9.3  # 9.3 度
    
    # 地平 30 度以上的天顶角
    min_altitude = 30 * u.deg
    
    # 使用 HEALPix 网格生成天区坐标
    nside = 32
    hp = HEALPix(nside=nside)
    ra, dec = hp.healpix_to_lonlat(np.arange(hp.npix))
    coords = SkyCoord(ra, dec, unit='deg')
    
    # 转换到 AltAz 坐标系
    frame = AltAz(obstime=target_time, location=obs_location)
    altaz = coords.transform_to(frame)
    
    # 筛选地平 30 度以上的区域
    visible = altaz.alt > min_altitude

    # 获取 SVOM 观测的中心
    svom_center_ra = None
    for start, end, center, _ in svom_obs:
        if start <= target_time.datetime <= end:
            svom_center_ra = center.ra.degree
            break
    if svom_center_ra is None:
        svom_center_ra = 0  # 默认中心

    # **使用 Cartopy 设置天球坐标**
    projection = ccrs.Mollweide(central_longitude=svom_center_ra)
    fig, ax = plt.subplots(figsize=(10, 6), subplot_kw={'projection': projection})

    # 绘制地平 30° 以上的天区覆盖范围
    ax.scatter(
        coords[visible].ra.wrap_at(180 * u.deg).degree,
        coords[visible].dec.degree,
        transform=ccrs.PlateCarree(),
        color='blue', s=1, alpha=0.5, label='Horizon > 30°'
    )

    # **绘制 SVOM 视场范围**
    for start, end, center, _ in svom_obs:
        if start <= target_time.datetime <= end:
            circle = patches.Circle(
                (center.ra.wrap_at(180 * u.deg).degree, center.dec.degree),
                radius=svom_field_radius,
                transform=ccrs.PlateCarree(),
                alpha=0.3, color='red', label='SVOM Field'
            )
            ax.add_patch(circle)
            break

    # **绘制 GWAC 视场范围**
    for tgObs in gwac_obs:
        time = tgObs['time']
        telescopes = tgObs['telescope_data']
        
        if time <= target_time.datetime <= time + timedelta(minutes=10):  # 假设每个观测持续 10 分钟
            for _, gwac_field in telescopes.items():
                circle = patches.Circle(
                    (gwac_field.ra.wrap_at(180 * u.deg).degree, gwac_field.dec.degree),
                    radius=gwac_field_radius,
                    transform=ccrs.PlateCarree(),
                    alpha=0.3, color='green', label='GWAC Field'
                )
                ax.add_patch(circle)
            break

    # **添加背景特征**
    ax.add_feature(cfeature.BORDERS, linestyle=':')
    ax.gridlines(draw_labels=True, color="gray", alpha=0.5)

    # **设置图例**
    handles, labels = ax.get_legend_handles_labels()
    by_label = dict(zip(labels, handles))  # 去重
    ax.legend(by_label.values(), by_label.keys(), loc='upper right')

    ax.set_title(f"Sky Coverage at {target_time.iso}", pad=20)
    plt.show()

def sampleOverlapData(data_list):
    # 解析时间并按时间排序
    for item in data_list:
        item['time'] = datetime.fromisoformat(item['time'])
    data_list.sort(key=lambda x: x['time'])
    
    # 筛选小于0.99的数据点，并分组
    under_099 = []
    temp_group = []
    for item in data_list:
        if item['total_overlap_area_10mount_sky'] < 0.99:
            if not temp_group or (item['time'] - temp_group[-1]['time'] <= timedelta(minutes=10)):
                temp_group.append(item)
            else:
                if len(temp_group) > 1:
                    duration = (temp_group[-1]['time'] - temp_group[0]['time']).total_seconds() / 60
                    if duration > 10:
                        under_099.extend(temp_group)
                temp_group = [item]
    
    # 处理最后一组
    if len(temp_group) > 1:
        duration = (temp_group[-1]['time'] - temp_group[0]['time']).total_seconds() / 60
        if duration > 10:
            under_099.extend(temp_group)
    
    # 筛选大于0.99的数据点
    over_099 = [item for item in data_list if item['total_overlap_area_10mount_sky'] >= 0.99]
    
    # 进行随机采样
    sampled_under_099 = random.sample(under_099, min(5, len(under_099)))
    sampled_over_099 = random.sample(over_099, min(5, len(over_099)))
    
    return sampled_under_099, sampled_over_099


def batch_plot_sky_coverage(gwac_obs, svom_obs, overlap_ratios, saveDir):
    
    imagePaths = []
    
    saveDir1 = f"{saveDir}/sky_coverage_plots"
    if not os.path.exists(saveDir1):
        os.makedirs(saveDir1)
        
    sampled_under_099, sampled_over_099 = sampleOverlapData(overlap_ratios)
    selData = sampled_under_099 + sampled_over_099
    
    for tdata in selData:
        #selTime = datetime(2025, 2, 27, 19, 32, tzinfo=timezone.utc)
        selTime = tdata['time']
        selTimeStr = selTime.strftime('%Y-%m-%dT%H:%M:%S')
        print(f"Plotting sky coverage at {selTimeStr}, overlap ratio: {tdata['total_overlap_area_10mount_sky']:.4f}")
    
        saveFile = f"{selTimeStr}.png"
        savePath = os.path.join(saveDir1, saveFile)
        plot_sky_coverage(gwac_obs, svom_obs, selTime, savePath)
        # break
        
        imagePaths.append(f"sky_coverage_plots/{saveFile}")
        
    return imagePaths